package com.star.langchain.learn.config;

import java.util.List;

import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AiConfig {


    /**
     * 智谱AI API密钥
     */
    private static final String API_KEY = "8418ff45af1d42b78ad2dd0809b140f0.5WGLzAYvWQwR0MHw";

    @Bean
    public ChatModel chatModel() {
        return OpenAiChatModel.builder()
                .baseUrl("https://open.bigmodel.cn/api/paas/v4/")
                .modelName("glm-4.5-flash").apiKey(API_KEY)
                .logRequests(true).logResponses(true)
                .maxRetries(1).build();
    }

    @Bean
    public StreamingChatModel streamingChatModel() {
        return OpenAiStreamingChatModel.builder()
                .baseUrl("https://open.bigmodel.cn/api/paas/v4/")
                .modelName("glm-4.5-flash").apiKey(API_KEY).logRequests(true)
                .logResponses(true)
                .returnThinking(true)
                .build();
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        return OpenAiEmbeddingModel.builder()
                .baseUrl("https://open.bigmodel.cn/api/paas/v4/")
                .modelName("embedding-3").apiKey(API_KEY).logRequests(true)
                .logResponses(true)
                .maxRetries(1).build();
    }

    @Bean
    public ToolSpecification calculator() {
        return ToolSpecification.builder().name("calculator").description("returns a sum of two numbers")
                .parameters(JsonObjectSchema.builder().addIntegerProperty("first").addIntegerProperty("second")
                        .required(List.of("first", "second")).build())
                .build();
    }

}
